# -*- coding: utf8 -*-
__author__ = 'Luoge'

# import some modules used in the example
from PyQt4 import QtCore, QtGui
import traceback
import time

from qgis.core import *
from osgeo import gdal, ogr, osr
from osgeo.gdalconst import *
import numpy

class Worker(QtCore.QObject):
    finished = QtCore.pyqtSignal(basestring)
    error = QtCore.pyqtSignal(Exception, basestring)
    progress = QtCore.pyqtSignal(float)

    '''Example worker for calculating the total area of all features in a layer'''
    def __init__(self, inputRasterName, outputRasterName, greenNo =2, redNo = 3, nirNo= 4, mrNo = 5):
        QtCore.QObject.__init__(self)
        # if isinstance(layer, QgsVectorLayer) is False:
        #     raise TypeError('Worker expected a QgsVectorLayer, got a {} instead'.format(type(layer)))
        # self.layer = layer
        self.killed = False

        self.Input_raster = inputRasterName
        self.Green_Band_NO = greenNo
        self.Red_Band_NO = redNo
        self.Near_Infrared_Band_NO = nirNo
        self.Middle_Infrared_Band_NO = mrNo
        self.Output_Raster_layer_name = outputRasterName

    def run(self):
        try:
            # Open and assign the contents of the raster file to a dataset
            dataset = gdal.Open(self.Input_raster, GA_ReadOnly)

            bandGreenArray = dataset.GetRasterBand(self.Green_Band_NO).ReadAsArray(0, 0, dataset.RasterXSize, dataset.RasterYSize).astype(numpy.float64)

            bandRedArray = dataset.GetRasterBand(self.Red_Band_NO).ReadAsArray(0, 0, dataset.RasterXSize, dataset.RasterYSize).astype(numpy.float64)

            bandNIRArray = dataset.GetRasterBand(self.Near_Infrared_Band_NO).ReadAsArray(0, 0, dataset.RasterXSize, dataset.RasterYSize).astype(numpy.float64)

            bandMIRArray = dataset.GetRasterBand(self.Middle_Infrared_Band_NO).ReadAsArray(0, 0, dataset.RasterXSize, dataset.RasterYSize).astype(numpy.float64)

            oneBandSize = dataset.RasterYSize*dataset.RasterXSize
            waterArray = numpy.ones(oneBandSize)

            step = oneBandSize // 100
            progress_count = 0
            for i in range(oneBandSize):
                if self.killed is True:
                    # kill request received, exit loop early
                    break
                if bandGreenArray.flat[i] + bandRedArray.flat[i] > bandNIRArray.flat[i]+bandMIRArray.flat[i]:
                    waterArray[i]=1
                else:
                    waterArray[i]=0
                # increment progress
                progress_count += 1
                if step == 0 or progress_count % step == 0:
                    self.progress.emit(int(progress_count * 100.0 / float(oneBandSize)))
            if self.killed is False:
	            self.progress.emit(100)

            waterArray.resize((dataset.RasterYSize,dataset.RasterXSize))

            # Calculate and save the resulting dataset
            driver = gdal.GetDriverByName("GTiff")
            outDataset = driver.Create(self.Output_Raster_layer_name,
                                       dataset.RasterXSize,
                                       dataset.RasterYSize,
                                       1,
                                       gdal.GDT_Int16)
            waterBand = outDataset.GetRasterBand(1)
            waterBand.WriteArray(waterArray)
            # Check if there is geotransformation or geoprojection
            # in the input raster and set them in the resulting dataset
            if dataset.GetGeoTransform() != None:
                outDataset.SetGeoTransform(dataset.GetGeoTransform())
            if dataset.GetProjection() != None:
                outDataset.SetProjection(dataset.GetProjection())
            dataset = None
            outDataset = None
        except Exception, e:
            # forward the exception upstream
            self.error.emit(e, traceback.format_exc())

        # emit finished signal
        self.finished.emit(self.Output_Raster_layer_name)

    def kill(self):
        self.killed = True